import os
import sys
sys.path.append(os.path.abspath(os.path.dirname(__file__) + '/' + '../..'))

import matplotlib as mpl

mpl.use('pgf')
pgf_with_latex = {                     
    "pgf.texsystem": "xelatex",        
    "pgf.rcfonts": False,
    "text.usetex": True,                
    "font.family": "Times New Roman",
    "pgf.preamble": [
        r"\usepackage{fontspec}",    
        r"\setmainfont{Times New Roman}",        
        r"\usepackage{unicode-math}",
        r"\setmathfont{XITS Math}"       
        ]
    }    
mpl.rcParams.update(pgf_with_latex)

import numpy as np
import matplotlib.pyplot as plt

from python.tools import (
    clean_folder
)

from python.numerical.functions import (
    calculate_bias_coefs
)

################
## Parameters ##
################

output_folder = './numerical/output/bias_coefficients'

params_DGP = {'rho': 0.80,
              'sigma_eps': 1.0} # Normalization
params_DGP['sigma_omega'] = 3.0 * params_DGP['sigma_eps']
L = 10 # Number of bias coefficients to plot

params_noisy = {'rho_perceived': params_DGP['rho'],
                'sigma_eps_perceived': params_DGP['sigma_eps'],
                'sigma_omega_perceived': params_DGP['sigma_omega'],
                'theta': 0.00}

params_noisy_diagnostic = {'rho_perceived': params_DGP['rho'],
                           'sigma_eps_perceived': params_DGP['sigma_eps'],
                           'sigma_omega_perceived': params_DGP['sigma_omega'],
                           'theta': 0.60}

params_noisy_overextrap = {'rho_perceived': 0.90,
                           'sigma_eps_perceived': params_DGP['sigma_eps'],
                           'sigma_omega_perceived': params_DGP['sigma_omega'],
                           'theta': 0.00}

params_noisy_underextrap = {'rho_perceived': 0.70,
                           'sigma_eps_perceived': params_DGP['sigma_eps'],
                           'sigma_omega_perceived': params_DGP['sigma_omega'],
                           'theta': 0.00}

###############
## Get plots ##
###############

clean_folder(output_folder)

# Calculate bias coefficients
coefs_noisy = calculate_bias_coefs(params_DGP, params_noisy, L)
coefs_noisy_diagnostic = calculate_bias_coefs(params_DGP, params_noisy_diagnostic, L)
coefs_noisy_overextrap = calculate_bias_coefs(params_DGP, params_noisy_overextrap, L)
coefs_noisy_underextrap = calculate_bias_coefs(params_DGP, params_noisy_underextrap, L)

# Save output
coefs_noisy.to_csv('{}/baseline.csv'.format(output_folder), index = False)
coefs_noisy_diagnostic.to_csv('{}/noisy_diagnostic.csv'.format(output_folder), index = False)
coefs_noisy_overextrap.to_csv('{}/noisy_overextrap.csv'.format(output_folder), index = False)
coefs_noisy_underextrap.to_csv('{}/noisy_underextrap.csv'.format(output_folder), index = False)

# Get graph
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, sharex=True, sharey=True)

ax1.plot(coefs_noisy['L'], coefs_noisy['idiosyncratic'], 'b--^', label = 'Idiosyncratic')
ax1.plot(coefs_noisy['L'], coefs_noisy['aggregate'], 'r-o', label = 'Aggregate')
ax1.plot(coefs_noisy['L'], coefs_noisy['composite'], 'y:v', label = 'Composite')
ax1.axhline(0, color = 'k', linestyle = '--', alpha = 0.50)
ax1.set_title('Baseline')
ax1.set_ylim(-0.8, 0.4)

ax2.plot(coefs_noisy_diagnostic['L'], coefs_noisy_diagnostic['idiosyncratic'], 'b--^', label = 'Idiosyncratic')
ax2.plot(coefs_noisy_diagnostic['L'], coefs_noisy_diagnostic['aggregate'], 'r-o', label = 'Aggregate')
ax2.plot(coefs_noisy_diagnostic['L'], coefs_noisy_diagnostic['composite'], 'y:v', label = 'Composite')
ax2.axhline(0, color = 'k', linestyle = '--', alpha = 0.50)
ax2.set_title('Diagnostic')

ax3.plot(coefs_noisy_overextrap['L'], coefs_noisy_overextrap['idiosyncratic'], 'b--^', label = 'Idiosyncratic')
ax3.plot(coefs_noisy_overextrap['L'], coefs_noisy_overextrap['aggregate'], 'r-o', label = 'Aggregate')
ax3.plot(coefs_noisy_overextrap['L'], coefs_noisy_overextrap['composite'], 'y:v', label = 'Composite')
ax3.axhline(0, color = 'k', linestyle = '--', alpha = 0.50)
ax3.set_title('Over-Extrapolation')
ax3.set_ylim(-0.8, 0.4)

ax4.plot(coefs_noisy_underextrap['L'], coefs_noisy_underextrap['idiosyncratic'], 'b--^', label = 'Idiosyncratic')
ax4.plot(coefs_noisy_underextrap['L'], coefs_noisy_underextrap['aggregate'], 'r-o', label = 'Aggregate')
ax4.plot(coefs_noisy_underextrap['L'], coefs_noisy_underextrap['composite'], 'y:v', label = 'Composite')
ax4.axhline(0, color = 'k', linestyle = '--', alpha = 0.50)
ax4.set_title('Under-Extrapolation')

handles, labels = ax4.get_legend_handles_labels()
lgd = fig.legend(handles, labels, bbox_to_anchor=(0.5, 0.00), ncol = 3, loc='upper center', frameon = False)

fig.tight_layout()

plt.savefig('{}/bias_coefficients.png'.format(output_folder), dpi = 600, bbox_extra_artists=(lgd,), bbox_inches='tight')
plt.savefig('{}/bias_coefficients.pgf'.format(output_folder), dpi = 600, bbox_extra_artists=(lgd,), bbox_inches='tight')